import os

import cv2
import numpy as np
import streamlit as st
import ShapeRecognition as Sr
import ImageCorrection as Ic
import LaneLineDetection as Lld

# '''
#     chat_input 是 streamlit 3.9+ 新增的专门用于AI聊天的聊天框
#     并且可以上传文件
# '''

if user_input := st.chat_input("来和我聊天吧",
                               accept_file=True,
                               file_type=['jpg', 'png', 'jpeg']):
    with st.chat_message("user"):
        if user_input and user_input.text:
            # st.markdown(user_input.text)
            st.markdown(user_input.text)

        if user_input and user_input["files"]:
            # st.image(user_input["files"][0])
            st.image(user_input["files"][0], width=300)

            # 建议将上传的图片保存在本地
            if not os.path.exists("./model_save_img"):
                os.mkdir("./model_save_img")
            file_path = os.path.join("./model_save_img", user_input["files"][0].name)

            # 读取上传的文件内容
            file_bytes = user_input["files"][0].read()
            # 将文件转换为numpy数组
            np_arr = np.frombuffer(file_bytes, np.uint8)
            # 使用OpenCV解码图像
            img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
            cv2.imwrite(file_path, img)
    with st.chat_message("assistant"):
        if user_input and user_input.text and user_input["files"]:
            if "灰度" in user_input.text:
                img = cv2.imread(file_path)
                img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
                st.markdown("灰度处理结果如下：")
                st.image(img, width=300)
            if "浮雕" in user_input.text or "滤镜" in user_input.text:
                img = cv2.imread(file_path)
                h, w, _ = img.shape
                img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
                img_FuDiao = np.zeros_like(img)
                for i in range(h):
                    for j in range(w - 2):
                        p0 = img_gray[i, j]
                        p2 = img_gray[i, j + 2]
                        newP = p0 - p2 + 150
                        if newP < 0:
                            newP = 0
                        elif newP > 255:
                            newP = 255
                        img_FuDiao[i, j] = newP
                st.markdown("浮雕处理结果如下：")
                st.image(img_FuDiao, width=300)
            if "二值" in user_input.text:
                img = cv2.imread(file_path)
                img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
                img_binary = cv2.threshold(img_gray, 150, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1]
                st.markdown("二值处理结果如下：")
                st.image(img_binary, width=300)
            if "旋转" in user_input.text:
                img = cv2.imread(file_path)
                h, w, _ = img.shape
                M = cv2.getRotationMatrix2D((w / 2, h / 2), 45, 1)
                img_rotate = cv2.warpAffine(img, M, (w, h), borderMode=cv2.BORDER_CONSTANT, borderValue=(255, 255, 255))
                st.markdown("旋转处理结果如下：")
                img_rotate = cv2.cvtColor(img_rotate, cv2.COLOR_BGR2RGB)
                st.image(img_rotate, width=300)
            if "边缘" in user_input.text:
                img = cv2.imread(file_path)
                img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
                img_binary = cv2.threshold(img_gray, 150, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1]
                img_gauss = cv2.GaussianBlur(img_binary, (5, 5), 0)
                img_edge = cv2.Canny(img_gauss, 50, 110)
                st.markdown("边缘处理结果如下：")
                st.image(img_edge, width=300)
            if "形状" in user_input.text or "识别" in user_input.text:
                img = Sr.shape_recognition(file_path)
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                st.markdown("形状识别结果如下：")
                st.image(img, width=600)
            if "矫正" in user_input.text:
                img = Ic.image_correction(file_path)
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                st.markdown("图片矫正结果如下：")
                st.image(img, width=400)
            if "车道" in user_input.text:
                frame = cv2.imread(file_path)
                frame = cv2.resize(frame, (0, 0), fx=0.3, fy=0.3)
                edges_img, height, width = Lld.img_preprocess(frame)
                st.image(user_input["files"][0], width=400)
                st.markdown("请输入掩码参数：")
                first_width = st.number_input("第一个点位置--width * ？", min_value=0.0, max_value=1.0, value=0.1,
                                              step=0.1)
                first_height = st.number_input("第一个点位置--height * ？", min_value=0.0, max_value=1.0, value=0.1,
                                               step=0.1)
                second_width = st.number_input("第二个点位置--width * ？", min_value=0.0, max_value=1.0, value=0.1,
                                               step=0.1)
                second_height = st.number_input("第二个点位置--height * ？", min_value=0.0, max_value=1.0, value=0.1,
                                                step=0.1)
                third_width = st.number_input("第三个点位置--width * ？", min_value=0.0, max_value=1.0, value=0.1,
                                              step=0.1)
                third_height = st.number_input("第三个点位置--height * ？", min_value=0.0, max_value=1.0, value=0.1,
                                               step=0.1)
                fourth_width = st.number_input("第四个点位置--width * ？", min_value=0.0, max_value=1.0, value=0.1,
                                               step=0.1)
                fourth_height = st.number_input("第四个点位置--height * ？", min_value=0.0, max_value=1.0, value=0.1,
                                                step=0.1)
                if "开始识别" in user_input.text:
                    ROI = np.array(
                        [[first_width * width, first_height * height], [second_width * width - 20, second_height * height],
                         [third_width * width + 20, third_height * height], [fourth_width * width, fourth_height * height]])
                    masked_edges = Lld.img_mask(edges_img, ROI)
                    img = Lld.line_detection(masked_edges)
                    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                    st.markdown("车道线识别结果如下：")
                    st.image(img, width=600)
